#!/usr/bin/python

import cx_Oracle
import pygraphviz as pgv
import plan_output as plan_parser
import string
import argparse
import sys
import sqlparse
import random

class PlanLine():
    """
    Class which contain the content of the plan table output
    """
    def __init__(self, plan_line, Plan):
        planLines           = Plan.plan_lines
        self.has_predicate  = '*' in plan_line[0]
        self.idOp           = int(string.replace(plan_line[0],'*',''))
        self.opIndent       = int(plan_line[1][0]) -1
        self.parentId       = 0

        if self.idOp > 1:
            if planLines[-1].opIndent == self.opIndent:
                self.parentId = planLines[-1].parentId
            elif planLines[-1].opIndent > self.opIndent:
                self.parentId = [ i.parentId for i in planLines if i.opIndent == self.opIndent ][0]
            else:
                self.parentId = planLines[-1].idOp

        self.operation      = plan_line[1][1:]
        self.name           = string.replace(plan_line[2][1:],' ','')
        
        if self.has_predicate:
            self.node           = ' '.join([ i for i in self.operation, self.name, str(Plan.predic[str(self.idOp)]) ])
        else:
            self.node           = ' '.join([ i for i in self.operation, self.name ])

        self.row            = int(plan_line[3]) 
        self.bytes          = int(plan_line[4]) 
        self.cost           = int(string.split(plan_line[5],'(')[0])
        self.time           = plan_line[6]
        Plan.total_cost     += self.cost           
        planLines.append(self)

class Plan():
    """
    Class which contain all informations about the execution plan
    """
    def __init__(self, plan_parse):
        self.plan_parse = plan_parse
        self.hash_value = plan_parse.hash_value
        self.query      = plan_parse.query  
        self.predic     = plan_parse.predic 
        self.project    = plan_parse.project
        self.total_cost = 0

        self.plan_lines = []
        [ PlanLine(plan_line, self) for plan_line in plan_parse.plan_lines ]

        self.graph      = pgv.AGraph(label='Total Cost:' + str(self.total_cost), bgcolor='white')


    def make_graph(self):
        """
        Build the graph representation
        """
        self.graph.add_edge(self.plan_lines[0].node + str(self.hash_value), self.plan_lines[1].node, color='red')
        for i in self.plan_lines[2:]:
            self.graph.add_edge(self.plan_lines[i.parentId].node,i.node, color='red')
        return self.graph

    def draw_graph(self):
        """
        Export the graph image 
        """
        self.make_graph()
        self.graph.layout('dot')
        self.graph.draw(str(self.hash_value) + '_plangraph.png')

    def print_lines(self):
        print [ p.node for p in self.plan_lines ]

    def copy(self):
        """
        Return a copy of the plan
        """
        plan = Plan(self.plan_parse)
        plan.hash_value = 'm' + str(plan.hash_value)
        return plan

    def merge(self, plan):
        self.make_graph()
        color = "#%2x%2x%2x" % (random.randint(0,256),random.randint(0,256),random.randint(0,256))

        self.graph.add_edge(plan.plan_lines[0].node + str(plan.hash_value), plan.plan_lines[1].node, color= color)

        self.total_cost += plan.plan_lines[0].cost 
        if not self.graph.has_node(plan.plan_lines[1].cost):
            self.total_cost += plan.plan_lines[1].cost

        for i in plan.plan_lines[2:]:
            if not self.graph.has_node(i.node):
                self.total_cost += i.cost
            self.graph.add_edge(plan.plan_lines[i.parentId].node,i.node, color= color)

##
parse = argparse.ArgumentParser()
parse.add_argument('-u','--uri', default='is11cs04/is11cs04@192.168.2.29/oracle', help='Oracle uri format: \'USER/PASSWORD@NET_SERVICE_NAME/POOLING_ID\'')
parse.add_argument('-m','--merge', action='store_true', default= False, help='')
group = parse.add_mutually_exclusive_group(required=True)
group.add_argument('-f', '--file', type=file, default=sys.stdin, help='File name. (default = stdin)')
group.add_argument('query',type=str, nargs='?', help='One or more query between \"\". Example: \"select * from a; select b from c;\"')
args = parse.parse_args()

uri = args.uri
if not args.query:
    querys = args.file.readlines()
else:
    querys = sqlparse.split(args.query)
##

connection = cx_Oracle.connect(uri)
cursor = connection.cursor()

dbms_xplan_display = 'SELECT * FROM table(dbms_xplan.display(\'PLAN_TABLE\',NULL,\'ALL\'))'

plans = []
totalcost = 0

for sql in querys:
    if not sql or sql == '\n' :
        continue
    cursor.execute('explain plan for ' + sql.strip(';\n'))
    cursor.execute(dbms_xplan_display)
    dbms_display_output = '\n'.join( [ str(i[0]) for i in cursor ] )

    print sql
    print dbms_display_output

    plans.append( Plan(plan_parser.datatable.parseString(dbms_display_output)) )
    if not args.merge:
        totalcost += plans[-1].total_cost
        plans[-1].draw_graph()
        print "Total Cost: %d" % totalcost

if args.merge:
    copy = plans[0].copy()
    for p in plans[1:]:
        copy.merge(p)
    copy.graph.graph_attr.update(label=str(copy.total_cost))
    copy.draw_graph()
